#include <memory>
#include <vector>
#include <chrono>
#include <map>
#include <string>
#include <cmath>

#include "rclcpp/rclcpp.hpp"
#include "moveit_msgs/msg/display_trajectory.hpp"
#include "std_msgs/msg/float32.hpp"
#include "std_msgs/msg/int32.hpp"

using namespace std::chrono_literals;

class TrajectoryPublisherNode : public rclcpp::Node
{
public:
  TrajectoryPublisherNode() : Node("trajectory_publisher_node")
  {
    // 订阅display_planned_path话题
    trajectory_subscriber_ = this->create_subscription<moveit_msgs::msg::DisplayTrajectory>(
      "/display_planned_path",
      10,
      std::bind(&TrajectoryPublisherNode::trajectoryCallback, this, std::placeholders::_1));
    
    // 订阅爪子话题
    zhuazi_left_subscriber_ = this->create_subscription<std_msgs::msg::Int32>(
      "/zhuazi_left",
      10,
      std::bind(&TrajectoryPublisherNode::zhuaziLeftCallback, this, std::placeholders::_1));
    
    zhuazi_right_subscriber_ = this->create_subscription<std_msgs::msg::Int32>(
      "/zhuazi_right",
      10,
      std::bind(&TrajectoryPublisherNode::zhuaziRightCallback, this, std::placeholders::_1));
    
    // 创建11个发布器，分别对应11个关节
    center_publisher_ = this->create_publisher<std_msgs::msg::Float32>("/center_control", 10);
    
    for (int i = 1; i <= 5; ++i)
    {
      std::string left_topic = "/arm_left" + std::to_string(i);
      std::string right_topic = "/arm_right" + std::to_string(i);
      
      left_publishers_.push_back(this->create_publisher<std_msgs::msg::Float32>(left_topic, 10));
      right_publishers_.push_back(this->create_publisher<std_msgs::msg::Float32>(right_topic, 10));
    }
    
    // 初始化所有关节的值为0.0
    last_joint_values_["xuan1"] = 0.0f;
    last_joint_values_["huakuai2"] = 0.0f;
    last_joint_values_["yuntai2"] = 0.0f;
    last_joint_values_["xuan12"] = 0.0f;
    last_joint_values_["zhuazi2"] = 0.0f;
    last_joint_values_["huakuai1"] = 0.0f;
    last_joint_values_["yuntai1"] = 0.0f;
    last_joint_values_["xuan11"] = 0.0f;
    last_joint_values_["zhuazi1"] = 0.0f;
    
    // 获取发布频率参数（每个轨迹点的发布间隔，单位：毫秒）
    this->declare_parameter("publish_interval_ms", 150);
    publish_interval_ms_ = this->get_parameter("publish_interval_ms").as_int();
    
    // 获取轨迹点采样间隔参数（每隔几个点发布一个，1表示发布所有点）
    this->declare_parameter("point_skip_interval", 4);
    point_skip_interval_ = this->get_parameter("point_skip_interval").as_int();
    if (point_skip_interval_ < 1) {
      point_skip_interval_ = 1;
    }
    
    // 获取末尾重复发布次数参数
    this->declare_parameter("end_point_repeat_count", 10);
    end_point_repeat_count_ = this->get_parameter("end_point_repeat_count").as_int();
    if (end_point_repeat_count_ < 0) {
      end_point_repeat_count_ = 0;
    }
    
    RCLCPP_INFO(this->get_logger(), "轨迹发布节点已启动");
    RCLCPP_INFO(this->get_logger(), "订阅话题: /display_planned_path, /zhuazi_left, /zhuazi_right");
    RCLCPP_INFO(this->get_logger(), "发布话题: /center, /arm_left1~5, /arm_right1~5");
    RCLCPP_INFO(this->get_logger(), "转发: /zhuazi_left -> /arm_left5, /zhuazi_right -> /arm_right5");
    RCLCPP_INFO(this->get_logger(), "发布间隔: %d ms", publish_interval_ms_);
    RCLCPP_INFO(this->get_logger(), "轨迹点采样间隔: 每 %d 个点发布一个", point_skip_interval_);
    RCLCPP_INFO(this->get_logger(), "末尾重复发布次数: %d 次", end_point_repeat_count_);
    RCLCPP_INFO(this->get_logger(), "注意: xuan1从弧度转角度后，90度映射为0度，再取相反数 (value = -(value*180/π - 90))");
    RCLCPP_INFO(this->get_logger(), "注意: zhuazi1, zhuazi2从弧度转角度后取负值 (value = -(value*180/π))");
    RCLCPP_INFO(this->get_logger(), "注意: xuan11, xuan12 将自动从弧度转换为角度");
    RCLCPP_INFO(this->get_logger(), "注意: huakuai1的值由0.1~0映射为0~-0.1, huakuai2的值由-0.1~0映射为0~-0.1");
    RCLCPP_INFO(this->get_logger(), "注意: yuntai1的值由0.0~-0.388映射为0.0~0.388");
  }

private:
  // 爪子左侧回调函数
  void zhuaziLeftCallback(const std_msgs::msg::Int32::SharedPtr msg)
  {
    auto float_msg = std_msgs::msg::Float32();
    float_msg.data = static_cast<float>(msg->data);
    
    RCLCPP_INFO(this->get_logger(), "收到 /zhuazi_left: %d，开始连续转发5次到 /arm_left5", 
                msg->data);
    
    // 连续转发5次，每次间隔100ms
    for (int i = 0; i < 5; ++i)
    {
      left_publishers_[4]->publish(float_msg);  // left_publishers_[4] 对应 /arm_left5
      RCLCPP_INFO(this->get_logger(), "转发 /arm_left5 第 %d/5 次: %.2f", 
                  i + 1, float_msg.data);
      
      // 如果不是最后一次，则等待100ms
      if (i < 4)
      {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
      }
    }
    
    RCLCPP_INFO(this->get_logger(), "/zhuazi_left 转发完成");
  }
  
  // 爪子右侧回调函数
  void zhuaziRightCallback(const std_msgs::msg::Int32::SharedPtr msg)
  {
    auto float_msg = std_msgs::msg::Float32();
    float_msg.data = static_cast<float>(msg->data);
    
    RCLCPP_INFO(this->get_logger(), "收到 /zhuazi_right: %d，开始连续转发5次到 /arm_right5", 
                msg->data);
    
    // 连续转发5次，每次间隔100ms
    for (int i = 0; i < 5; ++i)
    {
      right_publishers_[4]->publish(float_msg);  // right_publishers_[4] 对应 /arm_right5
      RCLCPP_INFO(this->get_logger(), "转发 /arm_right5 第 %d/5 次: %.2f", 
                  i + 1, float_msg.data);
      
      // 如果不是最后一次，则等待100ms
      if (i < 4)
      {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
      }
    }
    
    RCLCPP_INFO(this->get_logger(), "/zhuazi_right 转发完成");
  }
  
  void trajectoryCallback(const moveit_msgs::msg::DisplayTrajectory::SharedPtr msg)
  {
    RCLCPP_INFO(this->get_logger(), "收到轨迹数据，包含 %zu 条轨迹", 
                msg->trajectory.size());
    
    // 遍历所有轨迹（通常只有一条）
    for (size_t i = 0; i < msg->trajectory.size(); ++i)
    {
      const auto& robot_trajectory = msg->trajectory[i];
      
      // 检查是否有关节轨迹
      if (robot_trajectory.joint_trajectory.points.empty())
      {
        RCLCPP_WARN(this->get_logger(), "轨迹 %zu 为空，跳过", i);
        continue;
      }
      
      const auto& joint_names = robot_trajectory.joint_trajectory.joint_names;
      const auto& points = robot_trajectory.joint_trajectory.points;
      
      RCLCPP_INFO(this->get_logger(), "轨迹 %zu: 包含 %zu 个关节，%zu 个轨迹点", 
                  i, joint_names.size(), points.size());
      
      // 打印关节名称
      std::string joint_names_str = "关节名称: [";
      for (size_t j = 0; j < joint_names.size(); ++j)
      {
        joint_names_str += joint_names[j];
        if (j < joint_names.size() - 1)
          joint_names_str += ", ";
      }
      joint_names_str += "]";
      RCLCPP_INFO(this->get_logger(), "%s", joint_names_str.c_str());
      
      // 分次发布每个轨迹点
      publishTrajectoryPoints(points, joint_names);
    }
  }
  
  void publishTrajectoryPoints(
    const std::vector<trajectory_msgs::msg::JointTrajectoryPoint>& points,
    const std::vector<std::string>& joint_names)
  {
    // 创建关节名称到索引的映射
    std::map<std::string, size_t> joint_index_map;
    for (size_t j = 0; j < joint_names.size(); ++j)
    {
      joint_index_map[joint_names[j]] = j;
    }
    
    // 定义关节名称到发布器的映射关系
    // xuan1 -> center, huakuai2 -> left1, yuntai2 -> left2, xuan12 -> left3, zhuazi2 -> left4
    // huakuai1 -> right1, yuntai1 -> right2, xuan11 -> right3, zhuazi1 -> right4
    struct JointMapping {
      std::string joint_name;
      std::string topic_name;
      rclcpp::Publisher<std_msgs::msg::Float32>::SharedPtr publisher;
    };
    
    // 定义所有需要发布的关节映射
    std::vector<JointMapping> all_mappings = {
      {"xuan1", "/center", center_publisher_},
      {"huakuai2", "/arm_left1", left_publishers_[0]},
      {"yuntai2", "/arm_left2", left_publishers_[1]},
      {"xuan12", "/arm_left3", left_publishers_[2]},
      {"zhuazi2", "/arm_left4", left_publishers_[3]},
      {"huakuai1", "/arm_right1", right_publishers_[0]},
      {"yuntai1", "/arm_right2", right_publishers_[1]},
      {"xuan11", "/arm_right3", right_publishers_[2]},
      {"zhuazi1", "/arm_right4", right_publishers_[3]}
    };
    
    // 打印当前轨迹包含的关节
    std::string available_joints = "当前轨迹包含的关节: [";
    for (const auto& name : joint_names)
    {
      available_joints += name + " ";
    }
    available_joints += "]";
    RCLCPP_INFO(this->get_logger(), "%s", available_joints.c_str());
    
    // 构建完整的轨迹数据列表（包含所有关节的值）
    std::vector<std::map<std::string, float>> trajectory_data_list;
    
    for (size_t i = 0; i < points.size(); ++i)
    {
      const auto& point = points[i];
      
      // 检查轨迹点的位置数量是否与关节数量匹配
      if (point.positions.size() != joint_names.size())
      {
        RCLCPP_WARN(this->get_logger(), 
                    "轨迹点 %zu 的位置数量 %zu 与关节数量 %zu 不匹配，跳过该点",
                    i + 1, point.positions.size(), joint_names.size());
        continue;
      }
      
      // 创建当前轨迹点的完整数据（从上次的值开始）
      std::map<std::string, float> current_point_data = last_joint_values_;
      
      // 更新存在于当前轨迹中的关节值
      for (size_t j = 0; j < joint_names.size(); ++j)
      {
        const std::string& joint_name = joint_names[j];
        // 只更新我们关心的关节
        if (current_point_data.find(joint_name) != current_point_data.end())
        {
          float value = static_cast<float>(point.positions[j]);
          
          // xuan1特殊处理：弧度转角度后，90度映射为0度，再取相反数
          if (joint_name == "xuan1")
          {
            value = -(value * 180.0f / M_PI - 90.0f);
          }
          // zhuazi1和zhuazi2特殊处理：弧度转换为角度后取负值
          else if (joint_name == "zhuazi1" || joint_name == "zhuazi2")
          {
            value = -(value * 180.0f / M_PI);
          }
          // 其他旋转关节：弧度转换为角度
          else if (joint_name == "xuan11" || joint_name == "xuan12")
          {
            value = value * 180.0f / M_PI;
          }
          // huakuai1的值由0.1~0映射为0~-0.1
          else if (joint_name == "huakuai1")
          {
            // 输入范围：0.1~0，输出范围：0~-0.1
            value = value - 0.1f;
          }
          // huakuai2的值由-0.1~0映射为0~-0.1
          else if (joint_name == "huakuai2")
          {
            // 输入范围：-0.1~0，输出范围：0~-0.1
            value = -value - 0.1f;
          }
          // yuntai1的值由0.0到-0.388映射到0.0到0.388
          else if (joint_name == "yuntai1")
          {
            // 输入范围：0.0~-0.388，输出范围：0.0~0.388
            value = -value;
          }
          
          current_point_data[joint_name] = value;
        }
      }
      
      // 将这个完整的轨迹点数据添加到列表
      trajectory_data_list.push_back(current_point_data);
      
      // 更新last_joint_values_为当前点的值（为下一个点做准备）
      last_joint_values_ = current_point_data;
    }
    
    RCLCPP_INFO(this->get_logger(), "构建了包含 %zu 个轨迹点的完整数据列表", 
                trajectory_data_list.size());
    
    // 依次发布列表中的每个轨迹点（采样后）
    size_t published_count = 0;
    for (size_t i = 0; i < trajectory_data_list.size(); ++i)
    {
      // 采样策略：发布第一个点、最后一个点，以及按间隔采样的点
      bool should_publish = (i == 0) || // 第一个点
                           (i == trajectory_data_list.size() - 1) || // 最后一个点
                           (i % point_skip_interval_ == 0); // 按间隔采样
      
      if (!should_publish) {
        continue;
      }
      
      const auto& point_data = trajectory_data_list[i];
      published_count++;
      
      std::string log_msg = "发布轨迹点 " + std::to_string(i + 1) + "/" + 
                            std::to_string(trajectory_data_list.size()) + 
                            " (已发布 " + std::to_string(published_count) + " 个):";
      
      // 改进：交替发布左右臂数据以平衡下位机负载
      // 先发布中心关节
      {
        auto msg = std_msgs::msg::Float32();
        msg.data = point_data.at("xuan1");
        center_publisher_->publish(msg);
        log_msg += " xuan1=" + std::to_string(point_data.at("xuan1"));
        
        // 每个关节发布后都延迟5ms
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
      }
      
      // 交替发布左右臂关节，而不是连续发布同一侧
      for (size_t arm_idx = 0; arm_idx < 4; ++arm_idx)
      {
        // 先发布左臂第arm_idx个关节
        if (arm_idx < left_publishers_.size() - 1)  // 不包括arm_left5（爪子）
        {
          const std::string& left_joint = all_mappings[1 + arm_idx].joint_name;
          auto msg = std_msgs::msg::Float32();
          msg.data = point_data.at(left_joint);
          left_publishers_[arm_idx]->publish(msg);
          log_msg += " " + left_joint + "=" + std::to_string(point_data.at(left_joint));
          
          // 每个关节发布后都延迟5ms
          std::this_thread::sleep_for(std::chrono::milliseconds(5));
        }
        
        // 再发布右臂第arm_idx个关节
        if (arm_idx < right_publishers_.size() - 1)  // 不包括arm_right5（爪子）
        {
          const std::string& right_joint = all_mappings[5 + arm_idx].joint_name;
          auto msg = std_msgs::msg::Float32();
          msg.data = point_data.at(right_joint);
          right_publishers_[arm_idx]->publish(msg);
          log_msg += " " + right_joint + "=" + std::to_string(point_data.at(right_joint));
          
          // 每个关节发布后都延迟5ms
          std::this_thread::sleep_for(std::chrono::milliseconds(5));
        }
      }
      
      RCLCPP_INFO(this->get_logger(), "%s", log_msg.c_str());
      
      // 如果不是最后发布的点，等待一段时间再发布下一个点
      if (i < trajectory_data_list.size() - 1)
      {
        std::this_thread::sleep_for(std::chrono::milliseconds(publish_interval_ms_));
      }
    }
    
    RCLCPP_INFO(this->get_logger(), "轨迹点发布完成！总共 %zu 个点，实际发布 %zu 个点", 
                trajectory_data_list.size(), published_count);
    
    // 额外发布末尾数据点几次
    if (end_point_repeat_count_ > 0 && !trajectory_data_list.empty())
    {
      const auto& last_point_data = trajectory_data_list.back();
      RCLCPP_INFO(this->get_logger(), "开始额外发布末尾数据点 %d 次...", end_point_repeat_count_);
      
      for (int repeat = 1; repeat <= end_point_repeat_count_; ++repeat)
      {
        std::this_thread::sleep_for(std::chrono::milliseconds(200));  // 5Hz = 200ms间隔
        
        std::string log_msg = "额外发布末尾点 (" + std::to_string(repeat) + "/" + 
                              std::to_string(end_point_repeat_count_) + "):";
        
        // 改进：交替发布左右臂数据
        // 先发布中心关节
        {
          auto msg = std_msgs::msg::Float32();
          msg.data = last_point_data.at("xuan1");
          center_publisher_->publish(msg);
          log_msg += " xuan1=" + std::to_string(last_point_data.at("xuan1"));
          
          // 每个关节发布后都延迟5ms
          std::this_thread::sleep_for(std::chrono::milliseconds(5));
        }
        
        // 交替发布左右臂关节
        for (size_t arm_idx = 0; arm_idx < 4; ++arm_idx)
        {
          // 先发布左臂第arm_idx个关节
          if (arm_idx < left_publishers_.size() - 1)
          {
            const std::string& left_joint = all_mappings[1 + arm_idx].joint_name;
            auto msg = std_msgs::msg::Float32();
            msg.data = last_point_data.at(left_joint);
            left_publishers_[arm_idx]->publish(msg);
            log_msg += " " + left_joint + "=" + std::to_string(last_point_data.at(left_joint));
            
            // 每个关节发布后都延迟5ms
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
          }
          
          // 再发布右臂第arm_idx个关节
          if (arm_idx < right_publishers_.size() - 1)
          {
            const std::string& right_joint = all_mappings[5 + arm_idx].joint_name;
            auto msg = std_msgs::msg::Float32();
            msg.data = last_point_data.at(right_joint);
            right_publishers_[arm_idx]->publish(msg);
            log_msg += " " + right_joint + "=" + std::to_string(last_point_data.at(right_joint));
            
            // 每个关节发布后都延迟5ms
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
          }
        }
        
        RCLCPP_INFO(this->get_logger(), "%s", log_msg.c_str());
      }
      
      RCLCPP_INFO(this->get_logger(), "末尾数据点额外发布完成！");
    }
  }

  rclcpp::Subscription<moveit_msgs::msg::DisplayTrajectory>::SharedPtr trajectory_subscriber_;
  rclcpp::Subscription<std_msgs::msg::Int32>::SharedPtr zhuazi_left_subscriber_;
  rclcpp::Subscription<std_msgs::msg::Int32>::SharedPtr zhuazi_right_subscriber_;
  rclcpp::Publisher<std_msgs::msg::Float32>::SharedPtr center_publisher_;
  std::vector<rclcpp::Publisher<std_msgs::msg::Float32>::SharedPtr> left_publishers_;
  std::vector<rclcpp::Publisher<std_msgs::msg::Float32>::SharedPtr> right_publishers_;
  int publish_interval_ms_;
  int point_skip_interval_;  // 轨迹点采样间隔（每隔几个点发布一个）
  int end_point_repeat_count_;  // 末尾点重复发布次数
  
  // 保存每个关节的最后值，用于在轨迹不包含某些关节时保持其值不变
  std::map<std::string, float> last_joint_values_;
};

int main(int argc, char** argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<TrajectoryPublisherNode>();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}

